import os

def format_text(sample):
    text = f"<|im_start|>system\n你正在模仿吕小鱼的说话风格。<|im_end|>\n"
    text += f"<|im_start|>user\n{sample['input']}<|im_end|>\n"
    text += f"<|im_start|>assistant\n{sample['output']}<|im_end|>"
    return {"text": text}

from datasets import load_dataset

# 加载原始JSON文件
dataset = load_dataset("json", data_files="xiaoyu.jsonl", split="train")
dataset = dataset.map(format_text, remove_columns=["input", "output"])
################################
import torch
from transformers import (
    AutoModelForCausalLM, 
    AutoTokenizer,
    BitsAndBytesConfig,
    TrainingArguments,
)
from peft import LoraConfig, prepare_model_for_kbit_training,get_peft_model

# 4-bit量化配置（关键！）
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",        # 推荐的量化类型
    bnb_4bit_compute_dtype=torch.bfloat16,  # 计算时使用bfloat16
    bnb_4bit_use_double_quant=True    # 二次量化节省显存
)

# 加载模型与分词器
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen1.5-4B-Chat",
    quantization_config=bnb_config,
    device_map="auto",
    use_cache=False                   # 必须关闭缓存以兼容梯度检查点
)
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen1.5-4B-Chat")
# 安全设置 pad_token
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

# 准备4-bit训练
model = prepare_model_for_kbit_training(model)
################################
peft_config = LoraConfig(
    r=8,                             # Rank值（显存不足可降为4）
    lora_alpha=32,                   # 缩放因子
    target_modules=["q_proj", "v_proj"],  # 目标模块
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
    modules_to_save=["lm_head"]      # 保留输出层可训练
)

model = get_peft_model(model, peft_config)
################################
training_args = TrainingArguments(
    output_dir="./results",
    per_device_train_batch_size=1,   # 批大小必须为1
    gradient_accumulation_steps=8,   # 有效batch_size=8
    learning_rate=1e-4,              # 4-bit训练需更低学习率
    num_train_epochs=5,
    optim="paged_adamw_8bit",        # 使用分页优化器
    fp16=True,                       # 混合精度训练
    logging_steps=10,
    save_strategy="steps",
    save_steps=100,
    gradient_checkpointing=True,     # 梯度检查点（必须启用）
    report_to="none",
    max_grad_norm=0.3                # 防止梯度爆炸
)
################################
from trl import SFTTrainer,DataCollatorForCompletionOnlyLM

response_template = "<|im_start|>assistant\n"

# 创建数据对齐器（仅计算助手部分的损失）
collator = DataCollatorForCompletionOnlyLM(
    response_template=response_template,
    tokenizer=tokenizer,
    mlm=False
)

trainer = SFTTrainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    tokenizer=tokenizer,
    data_collator=collator,
)

# 开始训练
trainer.train()